#!/usr/bin/env python3
# @author Loganli
# 定时任务配置: */15 * * * * python3 /srv/loki/loki.py
# python3  /srv/loki/loki.py --loki-address "https://your-loki-server" --lark-webhook "https://your-webhook-url"

import requests
import argparse
from datetime import datetime, timedelta
import os

# 默认配置信息
DEFAULT_LOKI_ADDRESS = ""
DEFAULT_LARK_WEBHOOK_URL = ""

THRESHOLD_5XX = 500  # 5xx 告警阈值
THRESHOLD_4XX = 500  # 4xx 告警阈值
QUERY_RANGE = "15m"  # 查询时间范围

# 解析命令行参数
def parse_args():
    parser = argparse.ArgumentParser(description='Loki 监控告警脚本')
    parser.add_argument('--loki-address', 
                        default=os.environ.get('LOKI_ADDRESS', DEFAULT_LOKI_ADDRESS),
                        help='Loki 服务地址 (也可通过 LOKI_ADDRESS 环境变量设置)')
    parser.add_argument('--lark-webhook', 
                        default=os.environ.get('LARK_WEBHOOK_URL', DEFAULT_LARK_WEBHOOK_URL),
                        help='飞书 Webhook URL (也可通过 LARK_WEBHOOK_URL 环境变量设置)')
    return parser.parse_args()

def run_loki_query(query, loki_address):
    """使用 Loki HTTP API 执行范围查询"""
    url = f"{loki_address}/loki/api/v1/query_range"
    end = datetime.utcnow()
    start = end - timedelta(minutes=3)
    
    params = {
        "query": f'count_over_time({query}[{QUERY_RANGE}])',
        "start": int(start.timestamp()),
        "end": int(end.timestamp()),
        "step": "1s"
    }
    
    try:
        print(f"Executing query: {params['query']}")
        response = requests.get(url, params=params, timeout=30)
        response.raise_for_status()
        data = response.json()
        
        # 调试输出
        # print(data)
        
        return parse_query_result(data)
    except Exception as e:
        print(f"Loki query failed: {str(e)}")
        if hasattr(e, 'response') and e.response is not None:
            print(f"Response content: {e.response.text}")
        return 0

def parse_query_result(data):
    """解析 Loki 响应数据"""
    try:
        # 检查是否有结果
        if "status" not in data or data["status"] != "success":
            print(f"Query failed: {data.get('status', 'unknown')}")
            return 0
            
        if "data" not in data or "result" not in data["data"]:
            print("No results in response")
            return 0
            
        results = data["data"]["result"]
        if not results:
            print("Empty results")
            return 0
            
        # 对多个时间序列的结果求和
        total = 0
        for result in results:
            for value in result.get("values", []):
                try:
                    # value格式: [timestamp, count]
                    total += int(float(value[1]))  # 处理浮点数计数
                except (ValueError, TypeError):
                    continue
        
        print(f"Query returned {total} results")
        return total
        
    except Exception as e:
        print(f"Error parsing response: {str(e)}")
        return 0

def send_lark_alert(error_type, count, threshold, lark_webhook_url):
    """发送格式优化的飞书告警"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    title = f"ingress-nginx {error_type} 错误告警"
    
    # 将时间范围转换为可读格式
    range_minutes = int(QUERY_RANGE.rstrip('m'))
    time_range = f"{range_minutes}分钟"
    
    # 构建紧凑的告警内容
    content = "\n".join([
        f"**告警时间**: {timestamp}",
        f"**错误类型**: HTTP {error_type}",
        f"**当前数量**: {count} 条 ({time_range}内)",
        f"**触发阈值**: > {threshold} 条",
        f"**日志查询**: `{error_type}_query`"
    ])
    
    payload = {
        "msg_type": "interactive",
        "card": {
            "config": {"wide_screen_mode": True},
            "header": {
                "title": {"tag": "plain_text", "content": title},
                "template": "red" if "5" in error_type else "orange"
            },
            "elements": [
                {
                    "tag": "div",
                    "text": {
                        "tag": "lark_md",
                        "content": content
                    }
                },
                {
                    "tag": "note",
                    "elements": [
                        {
                            "tag": "plain_text",
                            "content": "Loki监控告警"
                        }
                    ]
                }
            ]
        }
    }
    
    try:
        resp = requests.post(lark_webhook_url, json=payload, timeout=10)
        print(f"Alert sent: {resp.status_code}")
        return True
    except Exception as e:
        print(f"Failed to send alert: {str(e)}")
        return False

if __name__ == "__main__":
    # 解析命令行参数
    args = parse_args()
    loki_address = args.loki_address
    lark_webhook_url = args.lark_webhook
    
    start_time = datetime.now()
    print(f"Starting Loki alert check at {start_time.strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"Loki address: {loki_address}")
    print(f"Query range: {QUERY_RANGE}")
    
    # 使用正确的查询语法
    query_5xx = r'{app="ingress-nginx"} | json | status =~ "5[0-9]{2}"'
    count_5xx = run_loki_query(query_5xx, loki_address)
    
    query_4xx = r'{app="ingress-nginx"} | json | status =~ "4[0-9]{2}"'
    count_4xx = run_loki_query(query_4xx, loki_address)
    
    print(f"5xx count: {count_5xx}, 4xx count: {count_4xx}")
    
    # 检查阈值并发送告警
    if count_5xx > THRESHOLD_5XX:
        print(f"Triggering 5xx alert ({count_5xx} > {THRESHOLD_5XX})")
        send_lark_alert("5xx", count_5xx, THRESHOLD_5XX, lark_webhook_url)
    
    if count_4xx > THRESHOLD_4XX:
        print(f"Triggering 4xx alert ({count_4xx} > {THRESHOLD_4XX})")
        send_lark_alert("4xx", count_4xx, THRESHOLD_4XX, lark_webhook_url)
    
    duration = (datetime.now() - start_time).total_seconds()
    print(f"Alert check completed in {duration:.2f} seconds")
